from cStringIO import StringIO
from ganges.datatypes import Container
from struct import Struct as _Struct



class ConstructError(Exception):
    pass
class SizeofError(ConstructError):
    pass
class FieldError(ConstructError):
    pass


class Construct(object):
    __slots__ = []
    def parse(self, data):
        return self.parse_stream(StringIO(data))
    def parse_stream(self, stream):
        return self._parse(None, stream, {})
    def build(self, obj):
        stream = StringIO()
        self.build_stream(obj, stream)
        return stream.getvalue()
    def build_stream(self, obj, stream):
        self._build(obj, stream, {})
    def sizeof(self, context = None):
        if context is None:
            context = {}
        return self._sizeof(context)
    def accept(self, visitor):
        return getattr(visitor, "visit_%s" % (self.__class__.__name__))(self)
    def _parse(self, obj, stream, context):
        raise NotImplementedError()
    def _build(self, obj, stream, context):
        raise NotImplementedError()
    def _sizeof(self, context):
        raise SizeofError("Can't calculate the size of %r" % (self,))

class Subconstruct(Construct):
    __slots__ = ["subcon"]
    def __init__(self, subcon):
        self.subcon = subcon
    def __repr__(self):
        return "%s(%r)" % (self.__class__.__name__, self.subcon)
    def _parse(self, obj, stream, context):
        return self.subcon._parse(obj, stream, context)
    def _build(self, obj, stream, context):
        return self.subcon._build(obj, stream, context)
    def _sizeof(self, context):
        return self.subcon._sizeof(context)

class StaticField(Construct):
    __slots__ = ["length"]
    def __init__(self, length):
        self.length = length
    def __repr__(self):
        return "%s(%r)" % (self.__class__.__name__, self.length)
    def _parse(self, obj, stream, context):
        data = stream.read(self.length)
        if len(data) != self.length:
            raise FieldError("expected %d, found %d" % (len(data), self.length))
        return data
    def _build(self, data, stream, context):
        if len(data) != self.length:
            raise FieldError("expected %d, found %d" % (len(data), self.length))
        stream.write(data)
    def _sizeof(self, context):
        return self.length

class FormattedField(StaticField):
    __slots__ = ["packer"]
    def __init__(self, endianity, formatchar):
        assert endianity in "!><"
        assert len(formatchar) == 1
        self.packer = _Struct(endianity + formatchar)
        StaticField.__init__(self, self.packer.size)
    def _parse(self, obj, stream, context):
        return self.packer.unpack(StaticField._parse(self, obj, stream, context))[0]
    def _build(self, value, stream, context):
        StaticField._build(self, self.packer.pack(value), stream, context)

class DynamicField(Construct):
    __slots__ = ["lengthfunc"]
    def __init__(self, lengthfunc):
        self.lengthfunc = lengthfunc
    def __repr__(self):
        return "%s(%r)" % (self.__class__.__name__, self.lengthfunc)
    def _parse(self, obj, stream, context):
        l = self.lengthfunc(context)
        data = stream.read(l)
        if len(data) != l:
            raise FieldError("expected %d, found %d" % (len(data), l))
        return data
    def _build(self, data, stream, context):
        l = self.lengthfunc(context)
        if len(data) != l:
            raise FieldError("expected %d, found %d" % (len(data), l))
        stream.write(data)
    def _sizeof(self, context):
        return self.lengthfunc(context)

class Struct(Construct):
    __slots__ = ["subcons", "container_factory", "embedded"]
    def __init__(self, *subcons, **kwargs):
        self.container_factory = kwargs.pop("container_factory", dict)
        self.embedded = kwargs.pop("embedded", False)
        if kwargs:
            raise TypeError("invalid keyword arguments %s" % (kwargs.keys(),))
        self.subcons = subcons
    def _parse(self, obj, stream, context):
        if not self.embedded:
            obj = self.container_factory()
        context2 = {"_" : context}
        for name, sc in self.subcons:
            if name is None:
                sobj = sc._parse(obj, stream, context2)
                obj[name] = sobj
                context2[name] = sobj
        return obj
    def _build(self, obj, stream, context):
        context2 = {"_" : context}
        for name, sc in self.subcons:
            if name is None:
                sobj = obj
            else:
                sobj = obj[name]
                context2[name] = sobj
            sc._build(sobj, stream, context2)
    def _sizeof(self, context):
        context2 = {"_" : context}
        return sum(sc._sizeof(context2) for name, sc in self.subcons)

class Embed(Subconstruct):
    __slots__ = []
    def __init__(self, subcon):
        Subconstruct.__init__(self, subcon)
        self.subcon.embedded = True
    def _parse(self, obj, stream, context):
        return self.subcon._parse(obj, stream, context["_"])
    def _build(self, obj, stream, context):
        return self.subcon._build(obj, stream, context["_"])

class Sequence(Construct):
    __slots__ = ["subcons", "container_factory"]
    def __init__(self, *subcons):
        self.container_factory = kwargs.pop("container_factory", list)
        if kwargs:
            raise TypeError("invalid keyword arguments %s" % (kwargs.keys(),))
        self.subcons = subcons
    def _parse(self, obj, stream, context):
        obj = self.container_factory()
        context2 = {"_" : context}
        for sc in self.subcons:
            sobj = sc._parse(stream, context2)
            obj.append(sobj)
        return obj
    def _build(self, obj, stream, context):
        context2 = {"_" : context}
        it = iter(obj)
        for sc in self.subcons:
            sobj = it.next()
            sc._build(sobj, stream, context2)
    def _sizeof(self, context):
        context2 = {"_" : context}
        return sum(sc._sizeof(context2) for sc in self.subcons)

class DynamicArray(Subconstruct):
    __slots__ = ["countfunc", "container_factory"]
    def __init__(self, countfunc, subcon):
        self.container_factory = kwargs.pop("container_factory", dict)
        if kwargs:
            raise TypeError("invalid keyword arguments %s" % (kwargs.keys(),))
        Subconstruct.__init__(self, subcon)
        self.countfunc = countfunc
    def _parse(self, obj, stream, context):
        count = self.countfunc(context)
        obj = []
        while count > 0:
            sobj = self.subcon._parse(stream, context)
            obj.append(sobj)
            count -= 1
        return obj
    def _build(self, obj, stream, context):
        count = self.countfunc(context)
        it = iter(obj)
        try:
            while count > 0:
                sobj = it.next()
                self.subcon._build(sobj, stream, context)
                count -= 1
        except StopIteration:
            pass
        return obj























